/*---------------------------------------------------------------------
                               头文件
-----------------------------------------------------------------------*/

// 用于唇语模式下，音视融合数据采集
// soc serverIp
// 保存路径  /data/vi  /data/ai

#include "rknn_api.h"
#include "client_t.h"
#include <stdio.h>
#include <errno.h>
#include <cstring>
#include <cstdlib>
#include <unistd.h>
#include "stdlib.h"
#include <pthread.h>
#include <signal.h>
#include <sys/poll.h>
#include <float.h>

#include "rk_defines.h"
#include "rk_debug.h"
#include "rk_mpi_ai.h"
#include "rk_mpi_ao.h"
#include "rk_mpi_aenc.h"
#include "rk_mpi_adec.h"
#include "rk_mpi_sys.h"
#include "rk_mpi_mb.h"
// #include </usr/include/eigen3/Eigen/Dense>
#include "test_comm_argparse.h"
#include <vector>
#include "taiic_key.h"
#include <fcntl.h>
#include "taiic_rknn.h"
#include "taiic_sensor.h"
#include "stft.h"

#include <thread>

using namespace std;
// using namespace Eigen;

#define SESSION 10000
#define TYPE 44
#define PICTURE_C_SAVE 50
#define PICTURE_C_CREATE 51

// RK_S32 sessionNum = 0;
// RK_S32 typeNUm = 0;

#define WIDTH 320
#define HEIGHT 320
#define PIXELS 1.5
#define AICOUNT 26

char aiCmd[256] = "/data/ai/";
char aiCmdFinal[256] = {0};

char viCmd[256] = "/data/vi/";
char viCmdFinal[256] = {0};

char aoPcm[256] = "/oem/usr/pcm/";
char aoPcmFinal[256] = {0};

int ai_save_count = 0;

FILE *aifp = NULL;
FILE *vifp = NULL;

int ai_save_flag = -1;

int ai_save = 0;
int vi_save = 0;

char buffer[1024];
int nCount;

tAiViSaveEndCallback *s_aivi_save_end = NULL;

char *serverAddr = "192.168.87.67"; // 默认的服务端ip 127.0.0.1 本地回环
// char *serverAddr = RK_NULL;
int vi_sock_cli;
// int ai_sock_cli;
#define BUF_SIZE 1024
/*---------------------------------------------------------------------
                            全局变量/宏定义/结构体/函数声明
-----------------------------------------------------------------------*/

static void sigterm_handler(int sig);

int ptt_flg = -1; // ptt按键控制，1表示默认，0表示按下ptt
int audio_flag = -1;
int cnt = 0;

// 判断当前图像是否正在采集
RK_BOOL is_model_run = RK_FALSE;
TAIIC_AI_AO_CTX_S *ctx_aio; // 音频采集、输出相关参数
TAIIC_VI_CTX_S *ctx_vi;     // 视频获取相关配置参数

static RK_S32 taiic_set_channel_params_ao(TEST_AO_CTX_S *params)
{
    AUDIO_DEV aoDevId = params->s32DevId;
    AO_CHN aoChn = params->s32ChnIndex;
    RK_S32 result = 0;
    AO_CHN_PARAM_S pstParams;
    memset(&pstParams, 0, sizeof(AO_CHN_PARAM_S));

    // 设置 AO 通道参数
    result = RK_MPI_AO_SetChnParams(aoDevId, aoChn, &pstParams);
    if (result != RK_SUCCESS)
    {
        RK_LOGE("ao set channel params, aoChn = %d", aoChn);
        return RK_FAILURE;
    }

    return RK_SUCCESS;
}

RK_S32 taiic_init_mpi_ao(TEST_AO_CTX_S *params)
{
    RK_S32 result;
    // 启用 AO 通道
    result = RK_MPI_AO_EnableChn(params->s32DevId, params->s32ChnIndex);
    if (result != 0)
    {
        RK_LOGE("ao enable channel fail, aoChn = %d, reason = %x", params->s32ChnIndex, result);
        return RK_FAILURE;
    }

    // set sample rate of input data
    // 启用 AO 重采样
    result = RK_MPI_AO_EnableReSmp(params->s32DevId, params->s32ChnIndex,
                                   (AUDIO_SAMPLE_RATE_E)params->s32ReSmpSampleRate);
    if (result != 0)
    {
        RK_LOGE("ao enable channel fail, reason = %x, aoChn = %d", result, params->s32ChnIndex);
        return RK_FAILURE;
    }

    return RK_SUCCESS;
}
void *sendDataThread(void *ptr)
{
    TEST_AO_CTX_S *params = reinterpret_cast<TEST_AO_CTX_S *>(ptr);
    // MB_POOL_CONFIG_S pool_config; //定义内存缓存池属性结构体
    // set default value for struct
    RK_U8 *srcData = RK_NULL;
    AUDIO_FRAME_S frame;
    RK_U64 timeStamp = 0;
    RK_S32 s32MilliSec = -1; //-1表示阻塞模式
    RK_S32 size = 0;
    RK_S32 result = 0;
    FILE *file = RK_NULL;
    RK_LOGI("params->s32ChnIndex : %d", params->s32ChnIndex);
    file = fopen(params->srcFilePath, "rb");

    if (file == RK_NULL)
    {
        RK_LOGE("open save file %s failed because %s.", params->srcFilePath, strerror(errno));
        goto __EXIT;
    }

    srcData = reinterpret_cast<RK_U8 *>(calloc(512, sizeof(RK_U8)));
    memset(srcData, 0, 512);
    while (1)
    {
        size = fread(srcData, 1, 512, file);
        // RK_PRINT("======fread pcm %d====\n", size);

        frame.u32Len = size;
        frame.u64TimeStamp = timeStamp++;
        frame.enBitWidth = find_bit_width(params->s32BitWidth);
        frame.enSoundMode = find_sound_mode(params->s32Channel);
        frame.bBypassMbBlk = RK_FALSE;

        MB_EXT_CONFIG_S extConfig;
        memset(&extConfig, 0, sizeof(extConfig));
        extConfig.pOpaque = srcData;
        extConfig.pu8VirAddr = srcData;
        extConfig.u64Size = size;
        // 创建一个内存缓存快
        RK_MPI_SYS_CreateMB(&(frame.pMbBlk), &extConfig);
    __RETRY:
        // 发送 AO 音频帧, -1表示阻塞模式
        result = RK_MPI_AO_SendFrame(params->s32DevId, params->s32ChnIndex, &frame, s32MilliSec);

        // RK_PRINT("=====ao send frame %d====\n",result);
        if (result < 0)
        {
            RK_LOGE("send frame fail, result = %d, TimeStamp = %lld, s32MilliSec = %d",
                    result, frame.u64TimeStamp, s32MilliSec);
            goto __RETRY;
        }
        // 释放一个已经获取的缓存块
        RK_MPI_MB_ReleaseMB(frame.pMbBlk);
        // RK_PRINT("======release MB=======\n");

        if (size <= 0)
        {
            RK_PRINT("========eof=======\n");
            break;
        }
    }

__EXIT:
    // 等待指定设备和通道播放完成
    RK_MPI_AO_WaitEos(params->s32DevId, params->s32ChnIndex, s32MilliSec);
    if (file)
    {
        fclose(file);
        file = RK_NULL;
    }
    free(srcData);
    return RK_NULL;
}
RK_S32 taiic_open_device_ao(TEST_AO_CTX_S *ctx)
{
    AUDIO_DEV aoDevId = ctx->s32DevId;
    AUDIO_SOUND_MODE_E soundMode;

    AIO_ATTR_S aoAttr;
    memset(&aoAttr, 0, sizeof(AIO_ATTR_S));

    if (ctx->chCardName)
    {
        snprintf(reinterpret_cast<char *>(aoAttr.u8CardName),
                 sizeof(aoAttr.u8CardName), "%s", ctx->chCardName);
    }

    aoAttr.soundCard.channels = ctx->s32DeviceChannel;
    aoAttr.soundCard.sampleRate = ctx->s32SampleRate;
    aoAttr.soundCard.bitWidth = AUDIO_BIT_WIDTH_16;

    AUDIO_BIT_WIDTH_E bitWidth = find_bit_width(ctx->s32BitWidth);
    if (bitWidth == AUDIO_BIT_WIDTH_BUTT)
    {
        goto __FAILED;
    }
    aoAttr.enBitwidth = bitWidth;
    aoAttr.enSamplerate = (AUDIO_SAMPLE_RATE_E)ctx->s32ReSmpSampleRate;
    soundMode = find_sound_mode(ctx->s32Channel);
    if (soundMode == AUDIO_SOUND_MODE_BUTT)
    {
        goto __FAILED;
    }
    aoAttr.enSoundmode = soundMode;
    aoAttr.u32FrmNum = ctx->s32PeriodCount;
    aoAttr.u32PtNumPerFrm = ctx->s32PeriodSize;

    aoAttr.u32EXFlag = 0;
    aoAttr.u32ChnCnt = 2;

    // 设置AO设备(声卡)参数
    RK_MPI_AO_SetPubAttr(aoDevId, &aoAttr);
    // 启用AO设备
    RK_MPI_AO_Enable(aoDevId);

    return RK_SUCCESS;
__FAILED:
    return RK_FAILURE;
}
RK_S32 taiic_deinit_mpi_ao(AUDIO_DEV aoDevId, AO_CHN aoChn)
{
    // 禁用 AO 重采样
    RK_S32 result = RK_MPI_AO_DisableReSmp(aoDevId, aoChn);
    if (result != 0)
    {
        RK_LOGE("ao disable resample fail, reason = %d", result);
        return RK_FAILURE;
    }
    // 禁用 AO 通道
    result = RK_MPI_AO_DisableChn(aoDevId, aoChn);
    if (result != 0)
    {
        RK_LOGE("ao disable channel fail, reason = %d", result);
        return RK_FAILURE;
    }

    return RK_SUCCESS;
}
RK_S32 taiic_close_device_ao(TEST_AO_CTX_S *ctx)
{
    AUDIO_DEV aoDevId = ctx->s32DevId;
    // 禁用AO设备
    RK_S32 result = RK_MPI_AO_Disable(aoDevId);
    if (result != 0)
    {
        RK_LOGE("ao disable fail, reason = %d", result);
        return RK_FAILURE;
    }
    return RK_SUCCESS;
}

RK_S32 lip_test_mpi_ao(TEST_AO_CTX_S *ctx, RK_U32 pcm_label)
{
    // RK_PRINT("======into ao ========\n");
    TEST_AO_CTX_S params;
    pthread_t tidSend;
    char tmpPath[256];
    if (taiic_open_device_ao(ctx) != RK_SUCCESS)
    {
        return RK_FAILURE;
    }

    memcpy(&(params), ctx, sizeof(TEST_AO_CTX_S));
    params.s32ChnIndex = 0;

    // snprintf(tmpPath, sizeof(tmpPath), "/data/pcm/%d.pcm",pcm_label);
    snprintf(tmpPath, sizeof(tmpPath), "/data/over.pcm");
    params.srcFilePath = tmpPath;
    RK_PRINT("======params path is %s============\n", params.srcFilePath);

    taiic_set_channel_params_ao(&params);
    taiic_init_mpi_ao(&params);
    pthread_create(&tidSend, RK_NULL, sendDataThread, reinterpret_cast<void *>(&params));

    pthread_join(tidSend, RK_NULL);
    taiic_deinit_mpi_ao(params.s32DevId, params.s32ChnIndex);

    taiic_close_device_ao(ctx);

    return RK_SUCCESS;
}

static void ai_vi_save_end(int ai, int vi)
{
    // printf("==ai is %d  vi is %d==\n", ai, vi);
    if (ai == 1 && vi == 1)
    {
        RK_LOGW("===soc start==\n");
        FILE *tmpViFp = fopen(viCmdFinal, "rb"); // 以二进制方式打开（创建）文件
        while ((nCount = fread(buffer, 1, 1024, tmpViFp)) > 0)
        {
            send(vi_sock_cli, buffer, nCount, 0);
            // printf("===send vi len is %d==\n", nCount);
        }
        fclose(tmpViFp);

        RK_LOGI("============trans========\n");

        sleep(0.1);
        // send(vi_sock_cli, aiCmdFinal, strlen(aiCmdFinal), 0); // 发送文件名
        FILE *tmpAiFp = fopen(aiCmdFinal, "rb"); // 以二进制方式打开（创建）文件
        while ((nCount = fread(buffer, 1, 1024, tmpAiFp)) > 0)
        {
            send(vi_sock_cli, buffer, nCount, 0);
            // printf("===send ai len is %d==\n", nCount);
        }
        RK_LOGW("===soc end==\n");
        sleep(0.3);
        system("buzzer_work 2");
        system("rm -rf /data/vi/*");
        system("rm -rf /data/ai/*");
        vi_save = 0;
        ai_save = 0;

        is_model_run = RK_FALSE;
        RK_LOGW("====please speak next session===\n");
    }
}

// static void *aiq_server(void *argv)
// {
//     system("rkaiq_3A_server &");
//     return ((void *)0);
// }

static void key_callback(TAIIC_KEY_STATE old_state, TAIIC_KEY_NAME key_name, TAIIC_KEY_STATE key_state)
{
    // 判断key_name去实现对应功能,key_name=0表示音频线模式
    RK_LOGE("old_state==%d, key_name==%d, key_state==%d,count=%d\n", old_state, key_name, key_state, cnt++);

    switch (key_name)
    {
    case TAIIC_PTT:

        if (key_state == 0) // 按键按下
        {
            if ((!is_model_run)) // 唇图模式
            {
                is_model_run = RK_TRUE;
                snprintf(aiCmdFinal, sizeof(aiCmdFinal), "%ssession%d_%d.pcm", aiCmd, sessionNum, typeNUm);
                RK_LOGI("=====ai cmd is %s===\n", aiCmdFinal);

                snprintf(viCmdFinal, sizeof(viCmdFinal), "%ssession%d_%d.YUV", viCmd, sessionNum, typeNUm);
                RK_LOGI("=====vi cmd is %s===\n", viCmdFinal);

                aifp = fopen(aiCmdFinal, "wb");
                vifp = fopen(viCmdFinal, "wb");

                // throwOut_thread(NULL, &aiq_server);
                
                system("echo 1 >/sys/class/pwm/pwmchip1/pwm0/enable"); // 开启补光灯
                system("echo b > /sys/devices/virtual/adw/adwdev/adwgpio");
                taiic_vi_get_release_frame_loop(ctx_vi);

                // system("echo 0 > /sys/class/gpio/gpio121/value");

                typeNUm++;
                if (typeNUm == TYPE)
                {
                    sessionNum++;
                    typeNUm = 0;
                }
                if (sessionNum == SESSION)
                {
                    sessionNum = 0;
                    exit(0);
                }
                // sleep(2);
            }
        }

        break;
    }
}

// data 一段音频的指针
// wlen 音频的长度
void ai_frame_data_cback(void *data, RK_U32 wlen)
{
    if (ai_save_flag == 1 && ai_save_count < AICOUNT) // 开始
    {
        // RK_LOGW("=====AI DATA LEN IS %d, count is %d===\n", wlen, ai_save_count);
        if (aifp)
        {
            ai_save_count++;
            fwrite(data, wlen, 1, aifp);
            fflush(aifp);
        }
    }
    else if (ai_save_count == AICOUNT) // 关闭
    {

        fclose(aifp);
        ai_save_flag = -1;
        ai_save_count = 0; // 控制语音启停
        ai_save = 1;       // ai数据保存完成的标志位，soc传输标志位
        RK_LOGW("=====save ai end==\n");

        s_aivi_save_end(ai_save, vi_save);
    }
    else
    {
        return;
    }
}

// data 一帧图像的指针
// wcount 当前图像的顺序
void vi_frame_data_cback(void *data, RK_U32 wcount)
{
    // RK_LOGW("====VI DATA IS %d==\n", wcount);

    if (wcount == (PICTURE_C_SAVE))
    {
        fclose(vifp);
        vi_save = 1; // vi数据保存完成的标志位，soc传输标志位
        RK_LOGW("=====save vi end===\n");
        s_aivi_save_end(ai_save, vi_save);
    }
    else if (wcount < PICTURE_C_SAVE)
    {
        if (wcount == 0)
        {
            ai_save_flag = 1; // 语音采集开始
        }

        if (vifp)
        {
            fwrite(data, WIDTH * HEIGHT * PIXELS, 1, vifp);
            fflush(vifp);
        }
    }
}

static int rknn_media_out_config_init(TEST_AO_CTX_S *ctx)
{
    ctx->srcFilePath = "/oem/usr/pcm/over.pcm";
    ctx->s32ChnNum = 1;
    ctx->s32SampleRate = 16000;
    ctx->s32ReSmpSampleRate = 16000;
    ctx->s32DeviceChannel = 2;
    ctx->s32Channel = 1;
    ctx->s32GetTrackMode = 10;
    ctx->s32BitWidth = 16;
    ctx->s32PeriodCount = 4;
    ctx->s32PeriodSize = 1024;
    ctx->chCardName = "hw:0,0";
    ctx->s32DevId = 0;
    RK_PRINT("========media out config finish=======\n");
    return 1;
}

void send_ai_vi_data()
{
    for (;;)
    {
        if (ai_save == 1 && vi_save == 1)
        {
            printf("===soc start==\n");
            FILE *tmpViFp = fopen(viCmdFinal, "rb"); // 以二进制方式打开（创建）文件
            while ((nCount = fread(buffer, 1, 1024, tmpViFp)) > 0)
            {
                send(vi_sock_cli, buffer, nCount, 0);
                // printf("===send vi len is %d==\n", nCount);
            }
            fclose(tmpViFp);
            system("rm -rf /data/vi/*");

            printf("============trans========\n");

            sleep(0.1);
            // send(vi_sock_cli, aiCmdFinal, strlen(aiCmdFinal), 0); // 发送文件名
            FILE *tmpAiFp = fopen(aiCmdFinal, "rb"); // 以二进制方式打开（创建）文件
            while ((nCount = fread(buffer, 1, 1024, tmpAiFp)) > 0)
            {
                send(vi_sock_cli, buffer, nCount, 0);
                // printf("===send ai len is %d==\n", nCount);
            }
            printf("===soc end==\n");
            system("rm -rf /data/ai/*");


            sleep(0.5);
            system("buzzer_work 2");
            vi_save = 0;
            ai_save = 0;
        }
    }
}

void ai_vi_save_end_callback(tAiViSaveEndCallback aivi_save_end)
{
    if (aivi_save_end != NULL)
    {
        s_aivi_save_end = aivi_save_end;
    }
}



/*----------------------------------------------------------------
                            主函数main
-----------------------------------------------------------------*/

int main(int argc, char **argv)
{
    // pthread_t tidSend;

    if (argc > 1)
    {
        // sessionNum = atoi(argv[1]); // session数
        // typeNUm = atoi(argv[2]);    // label数
        serverAddr = argv[1]; // 服务器的ip地址
    }
    // system("echo 0 > /sys/class/gpio/gpio121/value"); // mic->spk+
    printf("set server add is %s\n", serverAddr);

    system("mic0_taiic.sh");
    system("echo f > /sys/devices/virtual/adw/adwdev/adwgpio"); // mic_mic0

    // 初始化socket、连接服务器
    vi_sock_cli = init_socket(serverAddr);
    // ai_sock_cli = init_socket(serverAddr);

    DIR *myaidir = NULL;
    DIR *myvidir = NULL;
    if ((myaidir = opendir(aiCmd)) == NULL || (myvidir = opendir(viCmd)) == NULL) // 创建文件夹ai vi
    {
        bool ret1 = mkdir(aiCmd, 0755);
        bool ret2 = mkdir(viCmd, 0755);
        RK_LOGI("%s  and %s created sucess!/n", aiCmd, viCmd);
    }
    ai_vi_save_end_callback(ai_vi_save_end);

    // initial
    ctx_aio = reinterpret_cast<TAIIC_AI_AO_CTX_S *>(malloc(sizeof(TAIIC_AI_AO_CTX_S)));
    memset(ctx_aio, 0, sizeof(TAIIC_AI_AO_CTX_S));

    // 视频获取相关配置参数
    ctx_vi = reinterpret_cast<TAIIC_VI_CTX_S *>(malloc(sizeof(TAIIC_VI_CTX_S))); // 分配内存空间
    memset(ctx_vi, 0, sizeof(TAIIC_VI_CTX_S));                                   // 统一设置为0

    // params set
    RK_LOGI("------------params set------------\n");

    taiic_ai_ao_media_config_init(ctx_aio);
    ctx_aio->aio_trans.srcData = reinterpret_cast<RK_U8 *>(calloc(ctx_aio->ao_size, sizeof(RK_U8)));
    ctx_aio->s32SetTrackMode = 8; //

    snprintf(aoPcmFinal, sizeof(aoPcmFinal), "%s%d.pcm", aoPcm, typeNUm);
    ctx_aio->ao_ctx.srcFilePath = aoPcmFinal;

    // 配置 视频获取相关配置参数
    taiic_vi_media_config_init(ctx_vi);
    ctx_vi->width = WIDTH;
    ctx_vi->height = HEIGHT;
    ctx_vi->loopCountSet = PICTURE_C_CREATE;

    // 按键检测
    taiic_all_key_registers(ctx_aio, key_callback);

    RK_LOGI("====go to callback===\n");
    taiic_ai_frame_data_callback(ai_frame_data_cback); // 注册ai数据回调事件

    taiic_vi_frame_data_callback(vi_frame_data_cback); // 注册vi数据回调事件

    // setup the SIGINT to ctrl+c and handing de-init process
    signal(SIGINT, sigterm_handler);
    RK_MPI_SYS_Init();

    taiic_aio_data_capture(ctx_aio);

    Handle(ctx_aio, lipAOO1123);

    while (!ctx_aio->gAiExit)
    {
        sleep(0.3);
        if (ctx_aio->gAiExit)
        {
            system("killall soc");
        }
    }

__FAILED_AIO:
    if (ctx_aio)
    {
        free(ctx_aio);
        ctx_aio = RK_NULL;
    }
__FAILED_VI:
    if (ctx_vi)
    {
        free(ctx_vi);
        ctx_vi = RK_NULL;
    }

    RK_MPI_SYS_Exit();
    uint_socket(vi_sock_cli); // 关闭socket客户端
    // uint_socket(ai_sock_cli); // 关闭socket客户端
    return 0;
}
/*---------------------------------------------------------------
                          函数实现
-----------------------------------------------------------------*/

static void sigterm_handler(int sig)
{
    RK_LOGI("Catched SIGINT %d\n", sig);
    ctx_aio->gAiExit = RK_TRUE;
}
